import re
import datetime
from typing import Any, Dict, List, Optional, Union, Callable, Tuple

class SQLXType:
    """SQLX 数据类型基类"""
    def __init__(self, value: Any = None):
        self.value = value

    def __str__(self) -> str:
        return str(self.value)

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, SQLXType):
            return self.value == other.value
        return self.value == other

    def to_python(self) -> Any:
        return self.value

class Integer(SQLXType):
    def __init__(self, value: int = 0):
        super().__init__(int(value) if value is not None else 0)

class Float(SQLXType):
    def __init__(self, value: float = 0.0):
        super().__init__(float(value) if value is not None else 0.0)

class String(SQLXType):
    def __init__(self, value: str = ""):
        super().__init__(str(value) if value is not None else "")

class Boolean(SQLXType):
    def __init__(self, value: bool = False):
        super().__init__(bool(value) if value is not None else False)

class Date(SQLXType):
    def __init__(self, value: Union[str, datetime.date] = None):
        if isinstance(value, str):
            value = datetime.datetime.strptime(value, "%Y-%m-%d").date()
        super().__init__(value or datetime.date.today())

class Array(SQLXType):
    def __init__(self, value: List[Any] = None, element_type: type = SQLXType):
        super().__init__(value or [])
        self.element_type = element_type

    def append(self, item: Any):
        if not isinstance(item, self.element_type):
            item = self.element_type(item)
        self.value.append(item)

    def filter(self, func: Callable[[Any], bool]) -> 'Array':
        return Array([item for item in self.value if func(item)], self.element_type)

class Map(SQLXType):
    def __init__(self, value: Dict[str, Any] = None):
        super().__init__(value or {})

    def get(self, key: str, default: Any = None) -> Any:
        return self.value.get(key, default)

    def set(self, key: str, value: Any):
        self.value[key] = value

class JSON(SQLXType):
    def __init__(self, value: Union[str, Dict, List] = None):
        import json
        if isinstance(value, str):
            value = json.loads(value)
        super().__init__(value or {})

    def __str__(self) -> str:
        import json
        return json.dumps(self.value)

class Variable:
    """SQLX 变量"""
    def __init__(self, name: str, data_type: type, value: Any = None):
        self.name = name
        self.data_type = data_type
        self.value = data_type(value) if value is not None else data_type()

    def assign(self, value: Any):
        if not isinstance(value, self.data_type):
            value = self.data_type(value)
        self.value = value

class Environment:
    """SQLX 执行环境"""
    def __init__(self, parent: 'Environment' = None):
        self.parent = parent
        self.variables: Dict[str, Variable] = {}
        self.functions: Dict[str, Callable] = {}
        self.tables: Dict[str, 'Table'] = {}

    def declare_variable(self, name: str, data_type: type, value: Any = None) -> Variable:
        if name in self.variables:
            raise Exception(f"Variable '{name}' already exists")
        var = Variable(name, data_type, value)
        self.variables[name] = var
        return var

    def get_variable(self, name: str) -> Variable:
        if name in self.variables:
            return self.variables[name]
        if self.parent:
            return self.parent.get_variable(name)
        raise Exception(f"Variable '{name}' not found")

    def set_variable(self, name: str, value: Any):
        if name in self.variables:
            self.variables[name].assign(value)
            return
        if self.parent:
            self.parent.set_variable(name, value)
            return
        raise Exception(f"Variable '{name}' not found")

    def register_function(self, name: str, func: Callable):
        self.functions[name] = func

    def get_function(self, name: str) -> Callable:
        if name in self.functions:
            return self.functions[name]
        if self.parent:
            return self.parent.get_function(name)
        raise Exception(f"Function '{name}' not found")

    def create_table(self, name: str, columns: List[Dict[str, type]]) -> 'Table':
        table = Table(name, columns)
        self.tables[name] = table
        return table

    def get_table(self, name: str) -> 'Table':
        if name in self.tables:
            return self.tables[name]
        if self.parent:
            return self.parent.get_table(name)
        raise Exception(f"Table '{name}' not found")

class Table:
    """SQLX 表结构"""
    def __init__(self, name: str, columns: List[Dict[str, type]]):
        self.name = name
        self.columns = columns
        self.rows: List[Dict[str, SQLXType]] = []

    def insert(self, values: Dict[str, Any]):
        row = {}
        for col in self.columns:
            col_name = list(col.keys())[0]
            col_type = list(col.values())[0]
            if col_name in values:
                row[col_name] = col_type(values[col_name])
            else:
                row[col_name] = col_type()
        self.rows.append(row)

    def select(self, conditions: Callable[[Dict[str, SQLXType]], bool] = None) -> List[Dict[str, SQLXType]]:
        if conditions is None:
            return self.rows
        return [row for row in self.rows if conditions(row)]

class SQLXRuntime:
    """SQLX 运行时引擎"""
    def __init__(self):
        self.global_env = Environment()
        self._init_builtin_functions()
        self._init_builtin_types()

    def _init_builtin_functions(self):
        """初始化内置函数"""
        env = self.global_env

        # 数学函数
        env.register_function("ADD", lambda a, b: a + b)
        env.register_function("SUBTRACT", lambda a, b: a - b)
        env.register_function("MULTIPLY", lambda a, b: a * b)
        env.register_function("DIVIDE", lambda a, b: a / b if b != 0 else None)
        env.register_function("MOD", lambda a, b: a % b if b != 0 else None)

        # 字符串函数
        env.register_function("CONCAT", lambda a, b: str(a) + str(b))
        env.register_function("LENGTH", lambda s: len(str(s)))
        env.register_function("UPPER", lambda s: str(s).upper())
        env.register_function("LOWER", lambda s: str(s).lower())

        # 数组函数
        env.register_function("ARRAY_SUM", lambda arr: sum(item.value for item in arr.value) if arr.value else 0)
        env.register_function("ARRAY_LENGTH", lambda arr: len(arr.value))
        env.register_function("ARRAY_FILTER", lambda arr, func: arr.filter(func))

    def _init_builtin_types(self):
        """初始化内置类型"""
        env = self.global_env
        env.register_function("INT", lambda x: Integer(x))
        env.register_function("FLOAT", lambda x: Float(x))
        env.register_function("STRING", lambda x: String(x))
        env.register_function("BOOLEAN", lambda x: Boolean(x))
        env.register_function("DATE", lambda x: Date(x))
        env.register_function("ARRAY", lambda x: Array(x))
        env.register_function("MAP", lambda x: Map(x))
        env.register_function("JSON", lambda x: JSON(x))

    def execute(self, sqlx_code: str) -> Any:
        """执行 SQLX 代码"""
        # 解析 SQLX 代码并执行
        statements = self._parse_statements(sqlx_code)
        result = None
        
        for stmt in statements:
            result = self._execute_statement(stmt, self.global_env)
        
        return result

    def _parse_statements(self, sqlx_code: str) -> List[str]:
        """将 SQLX 代码解析为语句列表"""
        # 简单实现，基于分号分割语句，忽略引号内的分号
        statements = []
        current_stmt = ""
        in_quote = False
        quote_char = None

        for char in sqlx_code:
            if char in ["'", '"']:
                if not in_quote:
                    in_quote = True
                    quote_char = char
                elif quote_char == char:
                    in_quote = False
            
            if char == ';' and not in_quote:
                statements.append(current_stmt.strip())
                current_stmt = ""
            else:
                current_stmt += char
        
        if current_stmt.strip():
            statements.append(current_stmt.strip())
        
        return statements

    def _execute_statement(self, stmt: str, env: Environment) -> Any:
        """执行单个 SQLX 语句"""
        # 简化的解析器，处理主要的 SQLX 结构
        
        # 声明变量
        if stmt.upper().startswith("DECLARE "):
            return self._execute_declare(stmt, env)
        
        # 设置变量值
        elif stmt.upper().startswith("SET "):
            return self._execute_set(stmt, env)
        
        # IF 语句
        elif stmt.upper().startswith("IF "):
            return self._execute_if(stmt, env)
        
        # WHILE 循环
        elif stmt.upper().startswith("WHILE "):
            return self._execute_while(stmt, env)
        
        # FOR 循环
        elif stmt.upper().startswith("FOR "):
            return self._execute_for(stmt, env)
        
        # TRY-CATCH 结构
        elif stmt.upper().startswith("TRY "):
            return self._execute_try_catch(stmt, env)
        
        # CREATE TABLE 语句
        elif stmt.upper().startswith("CREATE TABLE "):
            return self._execute_create_table(stmt, env)
        
        # INSERT INTO 语句
        elif stmt.upper().startswith("INSERT INTO "):
            return self._execute_insert(stmt, env)
        
        # SELECT 语句
        elif stmt.upper().startswith("SELECT "):
            return self._execute_select(stmt, env)
        
        # 函数调用
        elif re.match(r'^\w+\s*\(', stmt):
            return self._execute_function_call(stmt, env)
        
        # 其他语句
        return self._execute_other(stmt, env)

    def _execute_declare(self, stmt: str, env: Environment) -> Variable:
        """执行 DECLARE 语句"""
        # 简化实现
        pattern = r'DECLARE\s+(\w+)\s+(\w+)(?:\s+DEFAULT\s+(.+))?'
        match = re.match(pattern, stmt, re.IGNORECASE)
        
        if not match:
            raise Exception(f"Invalid DECLARE statement: {stmt}")
        
        var_name, type_name, default_value = match.groups()
        
        # 映射类型名称到实际类型
        type_map = {
            'INT': Integer,
            'FLOAT': Float,
            'STRING': String,
            'BOOLEAN': Boolean,
            'DATE': Date,
            'ARRAY': Array,
            'MAP': Map,
            'JSON': JSON
        }
        
        if type_name not in type_map:
            raise Exception(f"Unknown type: {type_name}")
        
        data_type = type_map[type_name]
        
        # 计算默认值
        if default_value:
            # 简化处理，实际应该解析表达式
            value = self._evaluate_expression(default_value, env)
        else:
            value = None
        
        return env.declare_variable(var_name, data_type, value)

    def _execute_set(self, stmt: str, env: Environment) -> None:
        """执行 SET 语句"""
        # 简化实现
        pattern = r'SET\s+(\w+)\s*=\s*(.+)'
        match = re.match(pattern, stmt, re.IGNORECASE)
        
        if not match:
            raise Exception(f"Invalid SET statement: {stmt}")
        
        var_name, expression = match.groups()
        
        # 计算表达式值
        value = self._evaluate_expression(expression, env)
        
        # 设置变量值
        env.set_variable(var_name, value)

    def _evaluate_expression(self, expr: str, env: Environment) -> Any:
        """计算表达式值"""
        # 简化实现，实际需要完整的表达式解析器
        
        # 检查是否是字面量
        if re.match(r'^\d+$', expr):
            return Integer(int(expr))
        elif re.match(r'^\d+\.\d+$', expr):
            return Float(float(expr))
        elif re.match(r'^[\'"](.*)[\'"]$', expr):
            return String(expr[1:-1])
        elif expr.upper() == 'TRUE':
            return Boolean(True)
        elif expr.upper() == 'FALSE':
            return Boolean(False)
        
        # 检查是否是变量
        if expr in env.variables:
            return env.get_variable(expr).value
        
        # 检查是否是函数调用
        func_match = re.match(r'^(\w+)\s*\((.*)\)$', expr)
        if func_match:
            func_name, args_str = func_match.groups()
            
            # 解析参数
            args = []
            if args_str.strip():
                for arg in args_str.split(','):
                    arg = arg.strip()
                    args.append(self._evaluate_expression(arg, env))
            
            # 调用函数
            if func_name in env.functions:
                return env.get_function(func_name)(*args)
        
        # 简单表达式计算
        if '+' in expr:
            parts = expr.split('+', 1)
            left = self._evaluate_expression(parts[0].strip(), env)
            right = self._evaluate_expression(parts[1].strip(), env)
            return left + right
        
        # 更多表达式处理...
        
        raise Exception(f"Unsupported expression: {expr}")

    def _execute_if(self, stmt: str, env: Environment) -> Any:
        """执行 IF 语句"""
        # 简化实现，实际需要处理嵌套和 ELSEIF
        pattern = r'IF\s+(.+)\s+THEN\s+(.+)\s+ELSE\s+(.+)\s+END\s+IF'
        match = re.match(pattern, stmt, re.IGNORECASE | re.DOTALL)
        
        if not match:
            # 尝试匹配没有 ELSE 的 IF
            pattern = r'IF\s+(.+)\s+THEN\s+(.+)\s+END\s+IF'
            match = re.match(pattern, stmt, re.IGNORECASE | re.DOTALL)
            
            if not match:
                raise Exception(f"Invalid IF statement: {stmt}")
            
            condition, if_body = match.groups()
            else_body = ""
        else:
            condition, if_body, else_body = match.groups()
        
        # 计算条件表达式
        condition_value = self._evaluate_expression(condition, env)
        
        # 执行相应的代码块
        if condition_value:
            return self.execute(if_body)
        elif else_body:
            return self.execute(else_body)

    def _execute_while(self, stmt: str, env: Environment) -> None:
        """执行 WHILE 循环"""
        pattern = r'WHILE\s+(.+)\s+DO\s+(.+)\s+END\s+WHILE'
        match = re.match(pattern, stmt, re.IGNORECASE | re.DOTALL)
        
        if not match:
            raise Exception(f"Invalid WHILE statement: {stmt}")
        
        condition, body = match.groups()
        
        # 循环执行
        while self._evaluate_expression(condition, env):
            self.execute(body)

    def _execute_for(self, stmt: str, env: Environment) -> None:
        """执行 FOR 循环"""
        pattern = r'FOR\s+(\w+)\s+IN\s+(.+)\s+DO\s+(.+)\s+END\s+FOR'
        match = re.match(pattern, stmt, re.IGNORECASE | re.DOTALL)
        
        if not match:
            raise Exception(f"Invalid FOR statement: {stmt}")
        
        var_name, range_expr, body = match.groups()
        
        # 解析范围表达式
        range_parts = range_expr.split('..')
        if len(range_parts) != 2:
            raise Exception(f"Invalid range expression: {range_expr}")
        
        start_expr, end_expr = range_parts
        start = self._evaluate_expression(start_expr.strip(), env).value
        end = self._evaluate_expression(end_expr.strip(), env).value
        
        # 循环执行
        for i in range(start, end + 1):
            loop_env = Environment(parent=env)
            loop_env.declare_variable(var_name, Integer, i)
            self.execute(body)

    def _execute_try_catch(self, stmt: str, env: Environment) -> Any:
        """执行 TRY-CATCH 结构"""
        pattern = r'TRY\s+(.+)\s+CATCH\s*(?:\((.+)\))?\s+(.+)\s+END\s+TRY'
        match = re.match(pattern, stmt, re.IGNORECASE | re.DOTALL)
        
        if not match:
            raise Exception(f"Invalid TRY-CATCH statement: {stmt}")
        
        try_body, exception_type, catch_body = match.groups()
        
        try:
            return self.execute(try_body)
        except Exception as e:
            if not exception_type or isinstance(e, Exception):
                # 创建异常变量
                if exception_type:
                    exc_env = Environment(parent=env)
                    exc_env.declare_variable(exception_type, String, str(e))
                    return self.execute(catch_body)
                else:
                    return self.execute(catch_body)
            raise e

    def _execute_create_table(self, stmt: str, env: Environment) -> Table:
        """执行 CREATE TABLE 语句"""
        pattern = r'CREATE\s+TABLE\s+(\w+)\s*\((.+)\)'
        match = re.match(pattern, stmt, re.IGNORECASE)
        
        if not match:
            raise Exception(f"Invalid CREATE TABLE statement: {stmt}")
        
        table_name, columns_str = match.groups()
        
        # 解析列定义
        columns = []
        column_parts = re.split(r',\s*(?![^()]*\))', columns_str)
        
        for part in column_parts:
            part = part.strip()
            if not part:
                continue
                
            col_parts = part.split()
            if len(col_parts) < 2:
                raise Exception(f"Invalid column definition: {part}")
                
            col_name = col_parts[0]
            col_type = col_parts[1].upper()
            
            # 映射类型名称到实际类型
            type_map = {
                'INT': Integer,
                'FLOAT': Float,
                'VARCHAR': String,
                'TEXT': String,
                'BOOLEAN': Boolean,
                'DATE': Date
            }
            
            if col_type not in type_map:
                raise Exception(f"Unknown type: {col_type}")
            
            columns.append({col_name: type_map[col_type]})
        
        return env.create_table(table_name, columns)

    def _execute_insert(self, stmt: str, env: Environment) -> None:
        """执行 INSERT INTO 语句"""
        pattern = r'INSERT\s+INTO\s+(\w+)\s*\((.+)\)\s+VALUES\s*\((.+)\)'
        match = re.match(pattern, stmt, re.IGNORECASE)
        
        if not match:
            raise Exception(f"Invalid INSERT statement: {stmt}")
        
        table_name, columns_str, values_str = match.groups()
        
        # 获取表
        table = env.get_table(table_name)
        
        # 解析列名
        columns = [col.strip() for col in columns_str.split(',')]
        
        # 解析值
        values = []
        value_parts = re.split(r',\s*(?![^()]*\))', values_str)
        
        for part in value_parts:
            part = part.strip()
            values.append(self._evaluate_expression(part, env))
        
        # 创建插入数据字典
        insert_data = {}
        for col, val in zip(columns, values):
            insert_data[col] = val
        
        # 插入数据
        table.insert(insert_data)

    def _execute_select(self, stmt: str, env: Environment) -> List[Dict[str, SQLXType]]:
        """执行 SELECT 语句"""
        pattern = r'SELECT\s+(.+)\s+FROM\s+(\w+)(?:\s+WHERE\s+(.+))?'
        match = re.match(pattern, stmt, re.IGNORECASE)
        
        if not match:
            raise Exception(f"Invalid SELECT statement: {stmt}")
        
        columns_str, table_name, where_str = match.groups()
        
        # 获取表
        table = env.get_table(table_name)
        
        # 解析列
        columns = [col.strip() for col in columns_str.split(',')]
        
        # 解析条件
        if where_str:
            # 简化实现，实际需要完整的条件表达式解析
            def condition(row):
                # 这里应该解析 where_str 并根据行数据计算条件
                # 简化处理，总是返回 True
                return True
        else:
            condition = None
        
        # 执行查询
        result_rows = table.select(condition)
        
        # 构建结果集
        result = []
        for row in result_rows:
            result_row = {}
            for col in columns:
                if col == '*':
                    result_row.update(row)
                else:
                    result_row[col] = row.get(col, None)
            result.append(result_row)
        
        return result

    def _execute_function_call(self, stmt: str, env: Environment) -> Any:
        """执行函数调用"""
        func_match = re.match(r'^(\w+)\s*\((.*)\)$', stmt)
        if not func_match:
            raise Exception(f"Invalid function call: {stmt}")
        
        func_name, args_str = func_match.groups()
        
        # 解析参数
        args = []
        if args_str.strip():
            for arg in re.split(r',\s*(?![^()]*\))', args_str):
                arg = arg.strip()
                args.append(self._evaluate_expression(arg, env))
        
        # 调用函数
        if func_name in env.functions:
            return env.get_function(func_name)(*args)
        
        raise Exception(f"Function '{func_name}' not found")

    def _execute_other(self, stmt: str, env: Environment) -> Any:
        """执行其他类型的语句"""
        # 处理其他 SQLX 语句
        print(f"Unhandled statement: {stmt}")
        return None

# 示例使用
if __name__ == "__main__":
    runtime = SQLXRuntime()
    
    # 创建表
    create_table_code = """
    CREATE TABLE users (
        id INT,
        name VARCHAR(50),
        age INT,
        is_active BOOLEAN
    );
    """
    runtime.execute(create_table_code)
    
    # 插入数据
    insert_code = """
    INSERT INTO users (id, name, age, is_active) VALUES (1, 'Alice', 30, TRUE);
    INSERT INTO users (id, name, age, is_active) VALUES (2, 'Bob', 25, TRUE);
    INSERT INTO users (id, name, age, is_active) VALUES (3, 'Charlie', 35, FALSE);
    """
    runtime.execute(insert_code)
    
    # 查询数据
    select_code = "SELECT * FROM users WHERE age > 26;"
    result = runtime.execute(select_code)
    
    # 打印结果
    for row in result:
        print({k: v.value for k, v in row.items()})
    
    # 执行复杂 SQLX 代码
    complex_code = """
    DECLARE total_age INT DEFAULT 0;
    DECLARE user_count INT DEFAULT 0;
    DECLARE average_age FLOAT;
    
    FOR user IN (SELECT age FROM users WHERE is_active = TRUE) DO
        SET total_age = total_age + user.age;
        SET user_count = user_count + 1;
    END FOR;
    
    IF user_count > 0 THEN
        SET average_age = total_age / user_count;
        PRINT 'Average age of active users: ' || average_age;
    ELSE
        PRINT 'No active users found';
    END IF;
    """
    # 注意：上面的代码包含 PRINT 语句，在当前实现中会导致错误
    # 实际使用时需要添加对 PRINT 等语句的支持
    # runtime.execute(complex_code)    